package websocket

import (
	"fmt"

	"golang.org/x/net/websocket"
)

// WebSocketHelper结构体用于辅助操作WebSocket
type WebSocketHelper struct {
	Conn *websocket.Conn
}

// Connect用于建立WebSocket连接
func (w *WebSocketHelper) Connect(urlStr string) error {
	ws, err := websocket.Dial(urlStr, "", "http://localhost/")
	if err != nil {
		return fmt.Errorf("连接WebSocket失败: %v", err)
	}

	w.Conn = ws
	return nil
}

// SendMessage发送消息（这里简单按文本消息示例，可根据需求扩展类型支持）
func (w *WebSocketHelper) SendMessage(message string) error {
	if w.Conn == nil {
		return fmt.Errorf("SendMessage,尚未建立WebSocket连接")
	}
	err := websocket.Message.Send(w.Conn, message)
	if err != nil {
		return fmt.Errorf("发送消息失败: %v", err)
	}
	return nil
}

// ReceiveMessage接收消息（返回消息内容，可按需进一步处理返回类型等细节）
func (w *WebSocketHelper) ReceiveMessage() (string, error) {
	if w.Conn == nil {
		return "", fmt.Errorf("ReceiveMessage,尚未建立WebSocket连接")
	}
	var msg string
	err := websocket.Message.Receive(w.Conn, &msg)
	if err != nil {
		return "", fmt.Errorf("接收消息失败: %v", err)
	}
	return msg, nil
}

// Close关闭WebSocket连接
func (w *WebSocketHelper) Close() error {
	if w.Conn == nil {
		return fmt.Errorf("Close,尚未建立WebSocket连接")
	}
	return w.Conn.Close()
}
